import re

from docx import Document


def add_numbering_to_heading(heading, level, numbering_dict):
    # 获取当前级别的编号
    if level not in numbering_dict:
        numbering_dict[level] = 1
    else:
        numbering_dict[level] += 1

    # 重置子级别的编号
    for sub_level in range(level + 1, 10):
        if sub_level in numbering_dict:
            del numbering_dict[sub_level]

    # 构建编号字符串
    number_str = ''
    for lvl in range(1, level + 1):
        if lvl == 1:
            number_str += chinese_number(numbering_dict[lvl])
        elif lvl == 2:
            number_str += f".{numbering_dict[lvl]}"
        else:
            number_str += f".{numbering_dict[lvl]}"

    # 创建编号段落
    numbering_paragraph = Document().add_paragraph()
    numbering_run = numbering_paragraph.add_run(number_str + ' ')
    numbering_run.bold = True

    # 将编号段落插入到标题之前
    heading._element.addprevious(numbering_paragraph._element)


def chinese_number(num):
    chinese_numbers = ["零", "一", "二", "三", "四", "五", "六", "七", "八", "九"]
    units = ["", "十", "百", "千", "万", "亿"]
    result = ""

    if num == 0:
        return chinese_numbers[0]

    unit_index = 0
    while num > 0:
        digit = num % 10
        if digit != 0:
            result = chinese_numbers[digit] + units[unit_index] + result
        elif result and result[0] != chinese_numbers[0]:
            result = chinese_numbers[0] + result

        num //= 10
        unit_index += 1

    # 处理特殊情况
    result = re.sub(r'零+', '零', result)
    result = re.sub(r'零$', '', result)
    result = re.sub(r'^一十', '十', result)

    return result


def modify_headings(doc_path):
    doc = Document(doc_path)
    numbering_dict = {}

    for paragraph in doc.paragraphs:
        style = paragraph.style.name
        if style.startswith('Heading'):
            level = int(style[-1])
            add_numbering_to_heading(paragraph, level, numbering_dict)

    # 保存修改后的文档
    doc.save(doc_path + '_modified.docx')



if __name__ == "__main__":
    input_doc = 'D:\opt\zpj.docx'
    doc_path = input_doc
    modify_headings(doc_path)
